Skip to content

fix: af.ex.Analysis use_jax=True with bare af.Model fallback#1240

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/nautilus-jax-example
Apr 29, 2026
Merged

fix: af.ex.Analysis use_jax=True with bare af.Model fallback#1240
Jammy2211 merged 1 commit into
mainfrom
feature/nautilus-jax-example

Conversation

@Jammy2211

@Jammy2211 Jammy2211 commented Apr 29, 2026

Copy link
Copy Markdown
Collaborator

Summary

af.ex.Analysis.model_data_1d_from had two paths: a Collection branch that correctly forwarded xp=self._xp to each profile's model_data_from, and a bare-Model fallback (the except TypeError branch, hit when instance is a single non-iterable model component) that called instance.model_data_from(xvalues=xvalues) without xp. With use_jax=True on a bare af.Model(af.ex.Gaussian), this defaulted xp to numpy and triggered TracerArrayConversionError inside jax.jit. The Collection branch was already updated when _xp was wired in; the single-instance branch was missed. One-line fix: add , xp=self._xp.

Surfaced while building a JAX-jitted Nautilus example in autofit_workspace_test (separate PR to follow once this lands).

API Changes

None — bug fix in example Analysis class. af.ex.Analysis(use_jax=True) now works with a bare af.Model; previously it silently degraded to numpy and crashed under JIT.

Test Plan

  • python -m pytest test_autofit/ -x passes
  • Manual: af.Nautilus over af.Model(af.ex.Gaussian) with Analysis(use_jax=True) runs to completion (verified locally: log_Z ≈ -64.15; JAX: Applying vmap and jit to likelihood function log line confirms JIT path was exercised)

🤖 Generated with Claude Code

@claude
fix: pass xp to bare-Model fallback in af.ex.Analysis.model_data_1d_from
f5691e1 
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Apr 29, 2026
@Jammy2211 Jammy2211 merged commit 606614b into main Apr 29, 2026
3 checks passed
@Jammy2211

Jammy2211 commented Apr 29, 2026

Copy link
Copy Markdown
Collaborator Author

Smoke Test Results — 2026-04-29

Workspace Passed Failed Skipped Total
autofit_workspace 5 1 1 7
autofit_workspace_test 8 0 1 9
autogalaxy_workspace 6 0 0 6
autolens_workspace 7 0 0 7
autolens_workspace_test 11 0 1 12
euclid_strong_lens_modeling_pipeline 6 0 0 6
HowToLens 6 0 0 6
Failures (1 pre-existing, not caused by this PR)

autofit_workspace/scripts/overview/overview_1_the_basics.py

Pre-existing failure — confirmed by running with canonical main branch PyAutoFit (same error). The user-defined Gaussian class in this tutorial defines model_data_from(self, xvalues) without xp= kwarg, but af.ex.Analysis.model_data_1d_from now passes xp=self._xp (the fix in this PR). The tutorial's local Gaussian needs updating to accept **kwargs or the explicit xp param. Not a regression from this PR.

Traceback (most recent call last):
  File ".../autofit_workspace/scripts/overview/overview_1_the_basics.py", line 429, in <module>
    result = search.fit(model=model, analysis=analysis)
  File ".../autofit/non_linear/search/abstract_search.py", line 848, in _fit_bypass_test_mode
    analysis.log_likelihood_function(instance)
  File ".../autofit/example/analysis.py", line 70, in log_likelihood_function
    model_data_1d = self.model_data_1d_from(instance=instance)
  File ".../autofit/example/analysis.py", line 112, in model_data_1d_from
    model_data_1d += instance.model_data_from(xvalues=xvalues, xp=self._xp)
TypeError: Gaussian.model_data_from() got an unexpected keyword argument 'xp'
Skipped
  • autofit_workspace: searches/mcmc.py — matches mcmc no_run entry (Zeus section broken)
  • autofit_workspace_test: database/scrape/general.py — in no_run (NEEDS_FIX 2026-04-27)
  • autolens_workspace_test: database/scrape/general.py — in no_run (NEEDS_FIX 2026-04-27)

@Jammy2211 Jammy2211 mentioned this pull request Apr 29, 2026
Merged
2 tasks
@Jammy2211

Jammy2211 commented Apr 29, 2026

Copy link
Copy Markdown
Collaborator Author

Workspace PR: PyAutoLabs/autofit_workspace_test#17

This was referenced Apr 30, 2026
Merged
Merged
@Jammy2211 Jammy2211 deleted the feature/nautilus-jax-example branch April 30, 2026 14:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

pending-release PR queued for the next release build

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Jammy2211